import os
import base64
import wave
import pyaudio
import soundfile
import io
import requests
import base64
import json
import time
import rclpy
from rclpy.parameter import Parameter
from rclpy.node import Node
from std_msgs.msg import String

# 路径
PACKAGE_NAME = 'robot_voice'
ASSETS_DIR = os.path.join(os.path.dirname(
    os.path.abspath(__file__)), "../../../../share", PACKAGE_NAME)
WAVE_INTPUT_FILENAME = os.path.join(ASSETS_DIR, "input.wav")
# pyaudio参数
CHUNK = 1024


class TTS(Node):
    def __init__(self, name):
        super().__init__(name)
        self.get_logger().info("启动节点: %s" % name)

        # 服务器地址
        self.declare_parameter('host', '127.0.0.1')
        self.declare_parameter('port', 8090)

        self.get_logger().info("服务器地址: %s" % f'http://{self.get_parameter("host").value}:{self.get_parameter("port").value}')

        self.tts_subscribe = self.create_subscription(
            String, "tts", self.tts_callback, 10)

    def paddlespeech_request(self, url, data):
        """ 
        构造PaddleSpeech所需的post请求
        """
        res = requests.post(
            url=url,
            data=json.dumps(data)
        )

        if res.status_code == 200:
            res = res.json()
        else:
            self.get_logger().info("请求失败，错误代码: " + str(res.status_code))
            res = None
        return res

    def tts_callback(self, message):
        data = {
            "text": message.data,
            "spk_id": 0,
            "speed": 1.0,
            "volume": 1.0,
            "sample_rate": 0,
            # 这里先不保存，演示如何将base64编码还原成wav音频
            "save_path": None
        }

        url = f'http://{self.get_parameter("host").value}:{self.get_parameter("port").value}/paddlespeech/tts'
        res = self.paddlespeech_request(url, data)
        success = res['success']
        wav_base64 = res['result']['audio']

        if success:
            self.get_logger().info("TTS获取成功!")
            self.analysis(wav_base64)
            self.play()
        else:
            self.get_logger().info("TTS获取失败!")

    def analysis(self, wav_base64):
        # base64还原成wav
        audio_data_byte = base64.b64decode(wav_base64)
        # 重新读取sample
        samples, sample_rate = soundfile.read(
            io.BytesIO(audio_data_byte), dtype='float32')
        # 保存成audio
        soundfile.write(WAVE_INTPUT_FILENAME, samples, sample_rate)

    def play(self):
        f = wave.open(WAVE_INTPUT_FILENAME, "rb")
        p = pyaudio.PyAudio()

        stream = p.open(format=p.get_format_from_width(f.getsampwidth()),
                        channels=f.getnchannels(),
                        rate=f.getframerate(),
                        output=True)

        stream.write(f.readframes(f.getnframes()))  # 把数据写进流对象

        time.sleep(1)

        stream.stop_stream()
        stream.close()
        f.close()
        p.terminate()


def main(args=None):
    rclpy.init(args=args)
    node = TTS("text_to_speech")
    rclpy.spin(node)
    rclpy.shutdown()
